'''
Description: 
Author: Egrt
Date: 2022-10-01 13:00:26
LastEditors: Egrt
LastEditTime: 2022-10-01 15:55:18
'''
import torch
import torch.backends.cudnn as cudnn
import matplotlib.pyplot as plt
import numpy as np
import torch.utils.data as data
import torchvision.datasets as datasets
from utils.utils import cvtColor, preprocess_input, resize_image
from PIL import Image
from nets.arcface import Arcface
from utils.dataloader import LFWDataset
from utils.utils_metrics import test, plot_rocs
from sklearn.manifold import TSNE
from mpl_toolkits.mplot3d import axes3d
from tqdm import tqdm

def show_2D_tSNE(latent_vecs, target, title='t-SNE viz'):
    latent_vecs = latent_vecs
    latent_vecs_reduced = TSNE(n_components=2, random_state=0).fit_transform(latent_vecs)
    plt.scatter(latent_vecs_reduced[:, 0], latent_vecs_reduced[:, 1],
                c=target, cmap='jet')
    plt.colorbar()
    plt.show()
    plt.savefig(title+".jpg")


def show_3D_tSNE(latent_vecs, target, title='3D t-SNE viz'):
    latent_vecs = latent_vecs
    tsne = TSNE(n_components=3, random_state=0).fit_transform(latent_vecs)
    fig = plt.figure(figsize=(13,10))
    ax = fig.add_subplot(111, projection='3d')
    scatter = ax.scatter3D(tsne[:, 0], tsne[:, 1], tsne[:, 2], c=target, cmap='jet')
    ax.set_title(title)
    plt.colorbar(scatter)
    plt.show()
    plt.savefig(title+".jpg")

def test(test_loader, model, png_save_path, log_interval, batch_size, cuda):
    labels, embeds = [], []
    pbar = tqdm(enumerate(test_loader))
    for batch_idx, (image, label) in pbar:
        with torch.no_grad():
            #--------------------------------------#
            #   加载数据，设置成cuda
            #--------------------------------------#
            image      = image.type(torch.FloatTensor)
            if cuda:
                image  = image.cuda()
            #--------------------------------------#
            #   传入模型预测，获得预测结果
            #   获得预测结果的距离
            #--------------------------------------#
            embed     = model(image)

        #--------------------------------------#
        #   将结果添加进列表中
        #--------------------------------------#
        embeds.append(embed.data.cpu().numpy())
        labels.append(label.data.cpu().numpy())

    #--------------------------------------#
    #   转换成numpy
    #--------------------------------------#
    labels      = np.array([sublabel for label in labels for sublabel in label])
    embeds      = np.array([subembed for embed in embeds for subembed in embed])
    
    return embeds, labels

class FacenetDataset(data.Dataset):
    def __init__(self, input_shape, lines, random):
        self.input_shape    = input_shape
        self.lines          = lines
        self.random         = random

    def __len__(self):
        return len(self.lines)

    def rand(self, a=0, b=1):
        return np.random.rand()*(b-a) + a

    def __getitem__(self, index):
        annotation_path = self.lines[index].split(';')[1].split()[0]
        y               = int(self.lines[index].split(';')[0])

        image = cvtColor(Image.open(annotation_path))
        #------------------------------------------#
        #   翻转图像
        #------------------------------------------#
        if self.rand()<.5 and self.random: 
            image = image.transpose(Image.FLIP_LEFT_RIGHT)
        image = resize_image(image, [self.input_shape[1], self.input_shape[0]], letterbox_image = True)

        image = np.transpose(preprocess_input(np.array(image, dtype='float32')), (2, 0, 1))
        return image, y
        
if __name__ == "__main__":
    #--------------------------------------#
    #   是否使用Cuda
    #   没有GPU可以设置成False
    #--------------------------------------#
    cuda            = True
    #--------------------------------------#
    #   主干特征提取网络的选择
    #   mobilefacenet
    #   mobilenetv1
    #   iresnet18
    #   iresnet34
    #   iresnet50
    #   iresnet100
    #   iresnet200
    #--------------------------------------#
    backbone        = "mobilefacenet"
    #--------------------------------------#
    #   输入图像大小
    #--------------------------------------#
    input_shape     = [112, 112, 3]
    #--------------------------------------#
    #   训练好的权值文件
    #--------------------------------------#
    model_path      = "model_data/arcface_mobilefacenet.pth"
    #--------------------------------------#
    #   LFW评估数据集的文件路径
    #   以及对应的txt文件
    #--------------------------------------#
    annotation_path  = "t-sne2.txt"
    #--------------------------------------#
    #   评估的批次大小和记录间隔
    #--------------------------------------#
    batch_size      = 256
    log_interval    = 1
    #--------------------------------------#
    #   ROC图的保存路径
    #--------------------------------------#
    png_save_path   = "model_data/roc_test.png"
    with open(annotation_path,"r") as f:
        lines = f.readlines()
    test_loader = torch.utils.data.DataLoader(
        FacenetDataset(input_shape, lines, random = False))

    model = Arcface(backbone=backbone, mode="predict")

    print('Loading weights into state dict...')
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
    model  = model.eval()

    if cuda:
        model = torch.nn.DataParallel(model)
        cudnn.benchmark = True
        model = model.cuda()

    embeds, labels = test(test_loader, model, png_save_path, log_interval, batch_size, cuda)
    show_2D_tSNE(embeds, labels, title='ArcFace (t-SNE)')
        